import torch
from torch.utils.data import Dataset, DataLoader
import os
import h5py

# ----------------------------------------------------------------------------
# --- 5. DATASET 定义 (DATASET DEFINITIONS) ---
# ----------------------------------------------------------------------------

feature_dim = 1024 

class H5Dataset(Dataset):
    """
    用于训练和内部验证的数据集。
    - 'train' 模式: 采样 'num_patches_to_sample' 个 patch。
    - 'val' (或 'test') 模式: 加载所有 patch。
    """

    def __init__(
        self,
        feats_path,
        df,
        split,
        split_col,
        task_label_col,
        num_patches_to_sample=1024,
    ):
        self.df = df[df[split_col] == split]
        self.feats_path = feats_path
        self.num_patches_to_sample = num_patches_to_sample
        self.split = split
        self.feature_dim = feature_dim
        self.task_label_col = task_label_col

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        with h5py.File(
            os.path.join(self.feats_path, row["slide_id"] + ".h5"), "r"
        ) as f:
            features = torch.from_numpy(f["features"][:])  # 加载所有 [N_all, D]

        if self.split == "train":
            # 训练时的采样逻辑
            num_available = features.shape[0]
            if num_available >= self.num_patches_to_sample:
                # 不重复采样
                indices = torch.randperm(num_available)[: self.num_patches_to_sample]
            else:
                # 有重复采样
                indices = torch.randint(
                    num_available,
                    (self.num_patches_to_sample,),
                )
            features = features[indices]  # 采样后 [num_patches_to_sample, D]

        # 对于 'val', features 保持为 [N_all, D]

        label_val = float(row[self.task_label_col])  # 或 int(...) 如果是整数标签
        label = torch.tensor(label_val, dtype=torch.float32)
        return features, label


class ExternalH5Dataset(Dataset):
    """
    用于外部测试的数据集。
    - 总是加载所有 patch。
    """

    def __init__(self, feats_path, df, label_col="label"):
        self.df = df
        self.feats_path = feats_path
        self.label_col = label_col

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx, row_key=None):
        row = self.df.iloc[idx]

        if row_key is not None:
            slide_id = row[row_key]
        else:
            slide_id = row["slide_id"]
            # slide_id = row["case_id"]
        with h5py.File(os.path.join(self.feats_path, slide_id + ".h5"), "r") as f:
            features = torch.from_numpy(f["features"][:])  # [N_all, D]
        label = torch.tensor(row[self.label_col], dtype=torch.float32)
        return features, label
